import torch
import torch.nn as nn

from torcheeg.models import EEGNet, LSTM
from util import *

class MyEEGNet(nn.Module):
    def __init__(self, in_chs:int=100, n_electrodes:int=16, 
                 n_classes:int=3, **kwargs) -> None:
        super().__init__()
        self.eegnet = EEGNet(chunk_size=in_chs, num_electrodes=n_electrodes, 
                             num_classes=n_classes, **kwargs)
    
    def forward(self, x:torch.Tensor):
        x = x.unsqueeze(1)
        return self.eegnet(x)

class MyLSTM(nn.Module):
    def __init__(self, n_classes:int=3, **kwargs) -> None:
        super().__init__()
        self.lstm = LSTM(num_electrodes=16, num_classes=n_classes, **kwargs)
    
    def forward(self, x):
        return self.lstm(x)

class MyLSTM2(nn.Module):
    def __init__(self, in_chs:int=100, hid_chs:int=100,
                 dropout:float=0.5, n_class:int=3) -> None:
        super().__init__()
        self.conv = nn.Conv1d(in_channels=in_chs, out_channels=64, kernel_size=3, stride=1, padding=1)
        self.maxpooling = nn.MaxPool1d(kernel_size=2, stride=2)
        self.dropout = nn.Dropout(dropout)
        self.lstm = nn.LSTM(input_size=64, hidden_size=hid_chs, batch_first=True)
        self.linear1 = nn.Linear(in_features=hid_chs, out_features=50)
        self.linear2 = nn.Linear(in_features=50, out_features=n_class)
        self.act1 = nn.ReLU()
        self.act2 = nn.Softmax(dim=1)

    def forward(self, x:torch.Tensor):
        x = x.permute(0, 2, 1)
        x = self.conv(x)
        x = self.maxpooling(x)
        x = self.dropout(x)

        x = x.permute(0, 2, 1)
        x, _ = self.lstm(x)
        x = x[:, -1, :]
        
        x = self.linear1(x)
        x = self.act1(x)
        x = self.dropout(x)
        x = self.linear2(x)
        x = self.act2(x)
        return x

if __name__ == '__main__':
    import torchinfo
    model = MyLSTM()
    model.to('cuda')
    sample = torch.randn(size=(2, 16, 100), device='cuda')
    torchinfo.summary(model, sample.shape)
